#conda activate mmdet / OpenLabCluster
# python -m lilab.OpenLabCluster_train.a2b_iter_data_prepare $PROJECT_REPR $PROJECT_SEMISEQ2SEQ --epoch 3 --iter-from 0
#%%
import pickle
import os
import os.path as osp
import h5py
import numpy as np
import pickle
import pandas as pd
import argparse
import shutil
from lilab.openlabcluster_postprocess.s1_merge_3_file import get_assert_1_file
from lilab.OpenLabCluster_train.a1_mirror_mutual_filt_clippredpkl import factory_label_mirror_start0


parser = argparse.ArgumentParser()
parser.add_argument("dir_representitive", type=str)
parser.add_argument("dir_semiseq2seq", type=str)
parser.add_argument("--iter-from", type=int, default=0)
parser.add_argument("--epoch", type=int, default=3)
args = parser.parse_args()

dir_representitive = args.dir_representitive
dir_semiseq2seq = project = args.dir_semiseq2seq
iteri = args.iter_from
iterto = iteri+1
epoch = args.epoch

clippredpkl_old_perc66 = get_assert_1_file(osp.join(dir_representitive,'*.clippredpkl'))
predpkldata = pickle.load(open(clippredpkl_old_perc66, 'rb'))
ind_rawclip = predpkldata['ind_rawclip']
cluster_labels_sub = predpkldata['cluster_labels']
clipNames_sub = predpkldata['clipNames']
df_clipNames_sub = predpkldata['df_clipNames']
cluster_labels_sub = cluster_labels_sub - cluster_labels_sub.min() + 1 #start from 1
nK_mutual = predpkldata['nK_mutual']
nK_mirrorhalf = predpkldata['nK_mirrorhalf'] if 'nK_mirrorhalf' in predpkldata else predpkldata['nK_mirror_half']

_, fun_label_mirror = factory_label_mirror_start0(nK_mutual, nK_mirrorhalf)

print(dir_semiseq2seq + f'/output/semisupervise-decSeq-iter{iteri}-epoch{epoch}/*.clippredpkl')
predpklfile_iterstart = get_assert_1_file(dir_semiseq2seq + f'/output/semisupervise-decSeq-iter{iteri}-epoch{epoch}/*.clippredpkl')
predpkldata_iterstart = pickle.load(open(predpklfile_iterstart, 'rb'))
df_clipNames_iterstart = predpkldata_iterstart['df_clipNames']

df_clipNames_sub_iterstart = df_clipNames_iterstart.loc[ind_rawclip]
assert np.all(df_clipNames_sub_iterstart[['vnake', 'startFrame', 'isBlack']].values == df_clipNames_sub[['vnake', 'startFrame', 'isBlack']].values)

df_clipNames_iterstart['raw_index'] = np.arange(len(df_clipNames_iterstart))
df_clipNames_iterstart['cluster_labels'] = predpkldata_iterstart['cluster_labels']


#%% calculate new mirrors
df_sort = df_clipNames_iterstart.sort_values(by=['vnake', 'startFrame', 'isBlack']).copy()
df_sort.reset_index(inplace=True, drop=True)
ind_sort = df_sort.index.values
new_label_sort = df_sort['cluster_labels'] - 1 #start from 0
new_label_W = new_label_sort[::2]
new_label_B = new_label_sort[1::2]
new_label_B_mirror = fun_label_mirror(new_label_B)


ind_mirror = (new_label_W == new_label_B_mirror) & (new_label_W>=0)
print('mean mirror', np.mean(new_label_W == new_label_B_mirror))
indx_mirror_not = (np.where(~ind_mirror)[0][:,None] * 2 + [[0,1]]).ravel()

df_sort.loc[indx_mirror_not, 'cluster_labels'] = -1
df_clipNames_iterstartb = df_sort.sort_values(by=['raw_index'])
df_clipNames_iterstartb.reset_index(inplace=True, drop=True)
np.mean(df_clipNames_iterstartb.loc[ind_rawclip]['cluster_labels'].values == cluster_labels_sub)
df_clipNames_iterstartb.loc[ind_rawclip]['cluster_labels'] = cluster_labels_sub

assert np.all(df_clipNames_iterstart[['isBlack', 'vnake', 'startFrame']].values == df_clipNames_iterstartb[['isBlack', 'vnake', 'startFrame']].values)

#%% save data
cluster_labels = df_clipNames_iterstartb['cluster_labels'].values
cluster_labels[cluster_labels<0] = 0
print('new annot labels %.2f'%(np.mean(cluster_labels>0)))
np.save(osp.join(dir_semiseq2seq, f'label/label-iter{iterto}.npy'), cluster_labels)

shutil.copy(osp.join(dir_semiseq2seq, f'datasets/data.h5'),
            osp.join(dir_semiseq2seq, f'datasets/data-iter{iterto}.h5'))
with h5py.File(osp.join(dir_semiseq2seq, f'datasets/data-iter{iterto}.h5'), 'r+') as fp:
    if 'label' in fp: fp.pop('label')
    fp.create_dataset('label', data=cluster_labels)

print(f'annotation样本数量 {np.sum(cluster_labels>0)}/{len(cluster_labels)} = {np.sum(cluster_labels>0)/len(cluster_labels)*100:.2f}%' )
